import os
import tarfile
import urllib.request

from tqdm import tqdm

from config.config import DATA_RAW


class DownloadProgressBar(tqdm):
    """带进度条的文件下载器"""

    def update_to(self, b=1, bsize=1, tsize=None):
        if tsize is not None:
            self.total = tsize
        self.update(b * bsize - self.n)


def download_cifar10():
    """下载CIFAR-10数据集"""
    url = "https://www.cs.toronto.edu/~kriz/cifar-10-python.tar.gz"
    filename = url.split('/')[-1]
    save_path = DATA_RAW / filename

    os.makedirs(DATA_RAW, exist_ok=True)

    if not os.path.exists(save_path):
        print(f"Downloading {filename} to {save_path}")
        with DownloadProgressBar(unit='B', unit_scale=True, miniters=1) as t:
            urllib.request.urlretrieve(
                url,
                filename=save_path,
                reporthook=t.update_to
            )
        print("Download completed!")
    else:
        print("File already exists. Skipping download.")

    return save_path


def extract_tar(tar_path):
    """解压数据集"""
    extract_to = DATA_RAW / "cifar-10-batches-py"
    if not extract_to.exists():
        print("Extracting files...")
        with tarfile.open(tar_path, "r:gz") as tar:
            members = tar.getmembers()
            for member in tqdm(members, desc="Extracting"):
                tar.extract(member, path=DATA_RAW)
        print(f"Files extracted to {extract_to}")
    else:
        print("Files already extracted. Skipping extraction.")

    return extract_to
